use std::sync::{Arc, Mutex};

use anyhow::Error;
use faiss::Index;

use crate::{Engine, EngineBuilder};

pub struct FaissEngine {
    index: Arc<Mutex<faiss::index::IndexImpl>>,
}

pub struct InitConfig {
    path: String,
}

impl Default for InitConfig {
    fn default() -> Self {
        Self { path: ".".into() }
    }
}

impl FaissEngine {}

impl EngineBuilder for FaissEngine {
    fn init(config: &Self::InitConfig) -> anyhow::Result<Self>
    where
        Self: Sized,
    {
        let index = faiss::read_index(&config.path)?;
        Ok(FaissEngine {
            index: Arc::new(Mutex::new(index)),
        })
    }

    type InitConfig = InitConfig;

    fn build(self) -> anyhow::Result<impl crate::Engine> {
        Ok(self)
    }
}

impl Engine for FaissEngine {
    async fn search(
        &self,
        inputs: &[f32],
        contents: &crate::content::ContentStore,
        config: &crate::SearchConfig,
    ) -> anyhow::Result<crate::content::ContentCollection> {
        let mut index = self
            .index
            .try_lock()
            .map_err(|e| Error::msg(format!("Failed to get faiss index locked. Error : {}", e)))?;
        let index = index.search(inputs, config.max_result_length)?;
        let result = index
            .labels
            .iter()
            .map(|label| label.get().unwrap() as usize)
            .collect::<Vec<_>>();
        let result = contents.get_collection(&result);
        if let Some(result) = result {
            Ok(result)
        } else {
            Err(Error::msg("Failed to get serached item."))
        }
    }
}
